#include "mnist.h"

// 添加字节序转换函数
int reverse_int(int i) {
    unsigned char c1, c2, c3, c4;
    c1 = i & 255;
    c2 = (i >> 8) & 255;
    c3 = (i >> 16) & 255;
    c4 = (i >> 24) & 255;
    return ((int)c1 << 24) + ((int)c2 << 16) + ((int)c3 << 8) + c4;
}

// 读取MNIST数据集
void read_mnist_data(const char* image_path, const char* label_path, 
                    float **images, unsigned char **labels, int *num_items) {
    FILE *fp = fopen(image_path, "rb");
    if (!fp) {
        printf("Cannot open image file: %s\n", image_path);
        exit(1);
    }
    
    int magic_number, num_images, num_rows, num_cols;
    fread(&magic_number, sizeof(int), 1, fp);
    fread(num_items, sizeof(int), 1, fp);
    fread(&num_rows, sizeof(int), 1, fp);
    fread(&num_cols, sizeof(int), 1, fp);
    
    // 转换字节序
    magic_number = reverse_int(magic_number);
    *num_items = reverse_int(*num_items);
    num_rows = reverse_int(num_rows);
    num_cols = reverse_int(num_cols);
    
    printf("Magic number: %d\n", magic_number);
    printf("Number of images: %d\n", *num_items);
    printf("Image size: %dx%d\n", num_rows, num_cols);
    
    // 检查图像大小是否正确
    if (num_rows * num_cols != IMAGE_SIZE) {
        printf("Invalid image size: expected %d, got %d\n", IMAGE_SIZE, num_rows * num_cols);
        exit(1);
    }
    
    // 读取图像数据
    unsigned char *temp_images = malloc(*num_items * IMAGE_SIZE);
    fread(temp_images, 1, *num_items * IMAGE_SIZE, fp);
    fclose(fp);
    
    // 转换为float并归一化
    *images = malloc(*num_items * IMAGE_SIZE * sizeof(float));
    for(int i = 0; i < *num_items * IMAGE_SIZE; i++) {
        (*images)[i] = temp_images[i] / 255.0f;
    }
    free(temp_images);
    
    // 读取标签文件
    fp = fopen(label_path, "rb");
    if (!fp) {
        printf("Cannot open label file: %s\n", label_path);
        exit(1);
    }
    
    fread(&magic_number, sizeof(int), 1, fp);
    fread(&num_images, sizeof(int), 1, fp);
    
    // 转换标签文件的字节序
    magic_number = reverse_int(magic_number);
    num_images = reverse_int(num_images);
    
    if (num_images != *num_items) {
        printf("Number of labels does not match number of images\n");
        exit(1);
    }
    
    *labels = malloc(*num_items);
    fread(*labels, 1, *num_items, fp);
    fclose(fp);
}

// 创建神经网络
Network* create_network() {
    Network *net = malloc(sizeof(Network));
    
    net->input_weights = malloc(IMAGE_SIZE * HIDDEN_SIZE * sizeof(float));
    net->hidden_weights = malloc(HIDDEN_SIZE * NUM_CLASSES * sizeof(float));
    net->input_bias = malloc(HIDDEN_SIZE * sizeof(float));
    net->hidden_bias = malloc(NUM_CLASSES * sizeof(float));
    
    // 初始化权重和偏置
    for(int i = 0; i < IMAGE_SIZE * HIDDEN_SIZE; i++)
        net->input_weights[i] = (float)rand() / RAND_MAX - 0.5f;
    for(int i = 0; i < HIDDEN_SIZE * NUM_CLASSES; i++)
        net->hidden_weights[i] = (float)rand() / RAND_MAX - 0.5f;
        
    return net;
}

// 训练网络
void train_network(Network *net, float *images, unsigned char *labels, int num_items) {
    float learning_rate = 0.1f;
    
    for(int epoch = 0; epoch < 1; epoch++) {  // 从10改为3轮
        printf("Epoch %d/1 started...\n", epoch + 1);  // 修改显示信息
        int correct_in_epoch = 0;
        
        for(int i = 0; i < num_items; i++) {
            // 前向传播
            float hidden[HIDDEN_SIZE];
            float output[NUM_CLASSES];
            
            // 计算隐藏层
            for(int j = 0; j < HIDDEN_SIZE; j++) {
                hidden[j] = net->input_bias[j];
                for(int k = 0; k < IMAGE_SIZE; k++) {
                    hidden[j] += images[i * IMAGE_SIZE + k] * 
                                net->input_weights[k * HIDDEN_SIZE + j];
                }
                hidden[j] = 1.0f / (1.0f + expf(-hidden[j])); // sigmoid激活
            }
            
            // 计算输出层
            for(int j = 0; j < NUM_CLASSES; j++) {
                output[j] = net->hidden_bias[j];
                for(int k = 0; k < HIDDEN_SIZE; k++) {
                    output[j] += hidden[k] * net->hidden_weights[k * NUM_CLASSES + j];
                }
                output[j] = 1.0f / (1.0f + expf(-output[j]));
            }
            
            // 计算目标输出向量
            float target[NUM_CLASSES] = {0};
            target[labels[i]] = 1.0f;
            
            // 输出层的误差
            float output_delta[NUM_CLASSES];
            for(int j = 0; j < NUM_CLASSES; j++) {
                output_delta[j] = (output[j] - target[j]) * output[j] * (1 - output[j]);
            }
            
            // 隐藏层的误差
            float hidden_delta[HIDDEN_SIZE];
            for(int j = 0; j < HIDDEN_SIZE; j++) {
                hidden_delta[j] = 0;
                for(int k = 0; k < NUM_CLASSES; k++) {
                    hidden_delta[j] += output_delta[k] * net->hidden_weights[j * NUM_CLASSES + k];
                }
                hidden_delta[j] *= hidden[j] * (1 - hidden[j]);
            }
            
            // 更新权重和偏置
            // 更新隐藏层到输出层的权重
            for(int j = 0; j < HIDDEN_SIZE; j++) {
                for(int k = 0; k < NUM_CLASSES; k++) {
                    net->hidden_weights[j * NUM_CLASSES + k] -= 
                        learning_rate * output_delta[k] * hidden[j];
                }
            }
            
            // 更新输入层到隐藏层的权重
            for(int j = 0; j < IMAGE_SIZE; j++) {
                for(int k = 0; k < HIDDEN_SIZE; k++) {
                    net->input_weights[j * HIDDEN_SIZE + k] -= 
                        learning_rate * hidden_delta[k] * images[i * IMAGE_SIZE + j];
                }
            }
            
            // 更新偏置
            for(int j = 0; j < NUM_CLASSES; j++) {
                net->hidden_bias[j] -= learning_rate * output_delta[j];
            }
            for(int j = 0; j < HIDDEN_SIZE; j++) {
                net->input_bias[j] -= learning_rate * hidden_delta[j];
            }
            
            // 计算当前样本的预测结果
            int predicted = 0;
            float max_output = output[0];
            for(int j = 1; j < NUM_CLASSES; j++) {
                if(output[j] > max_output) {
                    max_output = output[j];
                    predicted = j;
                }
            }
            if(predicted == labels[i]) {
                correct_in_epoch++;
            }
            
            // 每处理1000张图片显示一次进度
            if((i + 1) % 1000 == 0) {
                printf("Processing image %d/%d, Current accuracy: %.2f%%\n", 
                       i + 1, num_items, 
                       100.0f * correct_in_epoch / (i + 1));
            }
        }
        
        printf("Epoch %d/1 completed, Accuracy: %.2f%%\n",  // 修改显示信息
               epoch + 1, 
               100.0f * correct_in_epoch / num_items);
    }
}

// 预测函数
int predict(Network *net, float *image) {
    float hidden[HIDDEN_SIZE];
    float output[NUM_CLASSES];
    
    // 前向传播
    for(int j = 0; j < HIDDEN_SIZE; j++) {
        hidden[j] = net->input_bias[j];
        for(int k = 0; k < IMAGE_SIZE; k++) {
            hidden[j] += image[k] * net->input_weights[k * HIDDEN_SIZE + j];
        }
        hidden[j] = 1.0f / (1.0f + expf(-hidden[j]));
    }
    
    for(int j = 0; j < NUM_CLASSES; j++) {
        output[j] = net->hidden_bias[j];
        for(int k = 0; k < HIDDEN_SIZE; k++) {
            output[j] += hidden[k] * net->hidden_weights[k * NUM_CLASSES + j];
        }
        output[j] = 1.0f / (1.0f + expf(-output[j]));
    }
    
    // 找出最大概率的类别
    int max_index = 0;
    for(int i = 1; i < NUM_CLASSES; i++) {
        if(output[i] > output[max_index]) {
            max_index = i;
        }
    }
    
    return max_index;
}

// 释放网络内存
void free_network(Network *net) {
    free(net->input_weights);
    free(net->hidden_weights);
    free(net->input_bias);
    free(net->hidden_bias);
    free(net);
}

// 打印图像函数
void print_image(float *image) {
    printf("\nInput Image:\n");
    for(int i = 0; i < 28; i++) {
        for(int j = 0; j < 28; j++) {
            float pixel = image[i * 28 + j];
            printf("%c", pixel > 0.5f ? '#' : '.');
        }
        printf("\n");
    }
}

// 打印预测概率
void print_probabilities(Network *net, float *image) {
    float hidden[HIDDEN_SIZE];
    float output[NUM_CLASSES];
    
    // 前向传播计算
    for(int j = 0; j < HIDDEN_SIZE; j++) {
        hidden[j] = net->input_bias[j];
        for(int k = 0; k < IMAGE_SIZE; k++) {
            hidden[j] += image[k] * net->input_weights[k * HIDDEN_SIZE + j];
        }
        hidden[j] = 1.0f / (1.0f + expf(-hidden[j]));
    }
    
    for(int j = 0; j < NUM_CLASSES; j++) {
        output[j] = net->hidden_bias[j];
        for(int k = 0; k < HIDDEN_SIZE; k++) {
            output[j] += hidden[k] * net->hidden_weights[k * NUM_CLASSES + j];
        }
        output[j] = 1.0f / (1.0f + expf(-output[j]));
    }
    
    // 打印每个数字的概率
    printf("\nPrediction probabilities:\n");
    for(int i = 0; i < NUM_CLASSES; i++) {
        printf("Digit %d: %.2f%%\n", i, output[i] * 100);
    }
    
    // 找出最可能的数字
    int predicted = 0;
    float max_prob = output[0];
    for(int i = 1; i < NUM_CLASSES; i++) {
        if(output[i] > max_prob) {
            max_prob = output[i];
            predicted = i;
        }
    }
    printf("\nMost likely digit: %d (%.2f%% confidence)\n", predicted, max_prob * 100);
}
